import os

from gxl_ai_utils.utils import utils_file

input_wav_path = "./scp_data/wenetspeech/wav.scp"
input_text_path = "./scp_data/wenetspeech/text"
output_manifest_dir = "./output/wenetspeech"

def do_split_scp_for_mk_k2_manifest(input_wav_path, input_text_path, output_manifest_dir,num_split):
    # 切分成100小块
    temp_dir = os.path.join(output_manifest_dir, "temp")
    utils_file.makedir_sil(temp_dir)
    wav_dict = utils_file.load_dict_from_scp(input_wav_path)
    text_dict = utils_file.load_dict_from_scp(input_text_path)
    wav_dict_list = utils_file.do_split_dict(list(wav_dict.values()), num_split)
    text_dict_list = []
    for wav_dict_i in wav_dict_list:
        text_dict_i = {}
        for key, value in text_dict.items():
            if key in wav_dict_i:
                text_dict_i[key] = value
        text_dict_list.append(text_dict_i)
    for i in range(100):
        temp_dir_i = os.path.join(temp_dir, f'temp_{i}')
        utils_file.makedir_sil(temp_dir_i)
        temp_wav_dict = wav_dict_list[i]
        temp_text_dict = text_dict_list[i]
        utils_file.write_dict_to_scp(temp_wav_dict, os.path.join(temp_dir_i, "wav.scp"))
        utils_file.write_dict_to_scp(temp_text_dict, os.path.join(temp_dir_i, "text"))

